from langchain.chains import LLMChain
from langchain.tools import Tool
from langchain.agents import AgentType, initialize_agent
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from mcp.client import MCPClient
import playwright.sync_api
import asyncio

# 1. 定义 Playwright 百度搜索工具 (MCP 工具实现)
def baidu_search(query: str) -> str:
    """使用 Playwright 在百度上搜索指定关键词"""
    with playwright.sync_api.sync_playwright() as p:
        browser = p.chromium.launch(headless=True)
        page = browser.new_page()
        page.goto("https://www.baidu.com")
        page.fill("input#kw", query)
        page.click("input#su")
        page.wait_for_selector("div.result.c-container", timeout=5000)
        
        # 提取前三结果
        results = page.query_selector_all("div.result.c-container")[:3]
        output = "\n".join([result.inner_text() for result in results])
        browser.close()
        return f"百度搜索 '{query}' 结果：\n{output}"

# 2. 封装 MCP 工具客户端
class MCPBaiduSearchTool:
    def __init__(self, server_url="http://localhost:8000"):
        self.client = MCPClient(server_address=server_url)
        self.client.connect()
    
    def run(self, query: str) -> str:
        """调用 MCP 工具执行百度搜索"""
        params = {"query": query}
        result = self.client.call_tool("baidu_search", params)
        return str(result)

# 3. 创建意图识别链
def create_intent_chain():
    """创建用户意图识别链"""
    prompt = PromptTemplate.from_template(
        "分析用户输入，判断意图类型：\n"
        "可选意图：['搜索', '其他']\n"
        "用户输入: {input}\n"
        "意图类型:"
    )
    return LLMChain(
        llm=ChatOpenAI(temperature=0),
        prompt=prompt,
        output_key="intent"
    )

# 4. 初始化代理和工具
def create_agent():
    """创建 LangChain 代理"""
    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
    
    # 创建工具集
    tools = [
        Tool(
            name="BaiduSearch",
            func=MCPBaiduSearchTool().run,
            description="在百度上搜索内容"
        )
    ]
    
    # 初始化带意图识别的代理
    return initialize_agent(
        tools,
        llm,
        agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
        verbose=True,
        handle_parsing_errors=True
    )

# 5. 主执行流程
def main():
    agent = create_agent()
    intent_chain = create_intent_chain()
    
    user_input = "帮我用百度查一下Python的最新特性"
    
    # 识别用户意图
    intent_result = intent_chain.invoke({"input": user_input})
    
    if "搜索" in intent_result["intent"]:
        # 提取搜索关键词
        keyword_prompt = PromptTemplate.from_template(
            "从以下用户输入中提取搜索关键词：\n{input}\n关键词:"
        )
        keyword_chain = LLMChain(
            llm=ChatOpenAI(temperature=0),
            prompt=keyword_prompt
        )
        keyword = keyword_chain.invoke({"input": user_input})["text"]
        
        # 执行搜索
        print(f"识别到搜索意图，关键词: {keyword}")
        result = agent.invoke({
            "input": f"使用BaiduSearch工具搜索: {keyword}"
        })
        print(result["output"])
    else:
        # 其他意图处理
        print("未识别到搜索意图，执行默认响应")
        response = ChatOpenAI().invoke(user_input)
        print(response.content)

# 运行示例
if __name__ == "__main__":
    main()